from torch.utils.data import DataLoader
from torchvision import datasets,transforms
from configs import TEST_BATCH_SIZE, TRAIN_BATCH_SIZE, IMG_SIZE

train_pipeline = transforms.Compose([
    transforms.RandomResizedCrop(IMG_SIZE),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
    transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))
])

test_pipeline = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(IMG_SIZE),
    transforms.ToTensor(),
    transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))
])

def prepare_train_data():
    train_data = datasets.CIFAR10(root='data/', train=True,transform=train_pipeline,download=True)
    train_loader = DataLoader(dataset = train_data,shuffle=True,batch_size=TRAIN_BATCH_SIZE)
    return train_loader

def prepare_test_data():
    test_data =datasets.CIFAR10(root='data/',train=False,transform=test_pipeline,download=True)
    test_loader = DataLoader(test_data,batch_size=TEST_BATCH_SIZE,shuffle=True)
    return test_loader
